#pragma once

#include <complex>
#include <mpi.h>
#include <mutex>
#include <type_traits>

namespace compi
{
namespace types
{

namespace detail
{
// RAII 管理复数派生 MPI_Datatype，避免重复构造与泄漏
struct ComplexDatatypeManager
{
    MPI_Datatype complex_float{MPI_DATATYPE_NULL};
    MPI_Datatype complex_double{MPI_DATATYPE_NULL};

    std::once_flag complex_float_once;
    std::once_flag complex_double_once;

    ~ComplexDatatypeManager()
    {
        // 兜底释放：只有在 MPI 已初始化且尚未 Finalize 的情况下才尝试释放
        int initialized = 0;
        MPI_Initialized(&initialized);
        int finalized = 0;
        MPI_Finalized(&finalized);
        if (initialized && !finalized)
        {
            // 通常应在 MPI_Finalize 之前调用 cleanup_before_finalize()
            cleanup_before_finalize();
        }
        // 若 MPI 未初始化或已 Finalize，跳过释放以避免未定义行为
    }

    void cleanup_before_finalize()
    {
        if (complex_float != MPI_DATATYPE_NULL)
        {
            MPI_Type_free(&complex_float);
            complex_float = MPI_DATATYPE_NULL;
        }
        if (complex_double != MPI_DATATYPE_NULL)
        {
            MPI_Type_free(&complex_double);
            complex_double = MPI_DATATYPE_NULL;
        }
    }

    static ComplexDatatypeManager & instance()
    {
        static ComplexDatatypeManager inst;
        return inst;
    }

    MPI_Datatype get_complex_float()
    {
        std::call_once(complex_float_once,
                       [this]
                       {
                           MPI_Type_contiguous(2, MPI_FLOAT, &complex_float);
                           MPI_Type_commit(&complex_float);
                       });
        return complex_float;
    }

    MPI_Datatype get_complex_double()
    {
        std::call_once(complex_double_once,
                       [this]
                       {
                           MPI_Type_contiguous(2, MPI_DOUBLE, &complex_double);
                           MPI_Type_commit(&complex_double);
                       });
        return complex_double;
    }
};
} // namespace detail

// 默认认为不是固定类型，后续对内建类型进行特化。
template <typename T>
struct is_fixed_type
{
    static constexpr bool value = false;
};

// 内建固定类型特化
template <>
struct is_fixed_type<char>
{
    static constexpr bool value = true;
};

template <>
struct is_fixed_type<signed char>
{
    static constexpr bool value = true;
};

template <>
struct is_fixed_type<unsigned char>
{
    static constexpr bool value = true;
};

template <>
struct is_fixed_type<short>
{
    static constexpr bool value = true;
};

template <>
struct is_fixed_type<unsigned short>
{
    static constexpr bool value = true;
};

template <>
struct is_fixed_type<int>
{
    static constexpr bool value = true;
};

template <>
struct is_fixed_type<unsigned int>
{
    static constexpr bool value = true;
};

template <>
struct is_fixed_type<long>
{
    static constexpr bool value = true;
};

template <>
struct is_fixed_type<unsigned long>
{
    static constexpr bool value = true;
};

template <>
struct is_fixed_type<long long>
{
    static constexpr bool value = true;
};

template <>
struct is_fixed_type<unsigned long long>
{
    static constexpr bool value = true;
};

template <>
struct is_fixed_type<float>
{
    static constexpr bool value = true;
};

template <>
struct is_fixed_type<double>
{
    static constexpr bool value = true;
};

template <>
struct is_fixed_type<long double>
{
    static constexpr bool value = true;
};

template <>
struct is_fixed_type<bool>
{
    static constexpr bool value = true;
};

// std::byte 视为固定字节类型（单元素传输），允许通用字节通道
template <>
struct is_fixed_type<std::byte>
{
    static constexpr bool value = true;
};

// 复数类型：通过派生类型映射为固定类型
template <>
struct is_fixed_type<std::complex<float>>
{
    static constexpr bool value = true;
};

template <>
struct is_fixed_type<std::complex<double>>
{
    static constexpr bool value = true;
};

// 便捷变量模板
template <typename T>
inline constexpr bool is_fixed_type_v = is_fixed_type<T>::value;

// 为固定类型提供 MPI_Datatype 映射函数
// 参数：无（模板参数 T 表示要映射的类型）
// 返回值：对应的 MPI_Datatype；如果类型未被识别为固定类型，将返回 MPI_DATATYPE_NULL。
// 异常：无；调用者应确保仅在 is_fixed_type_v<T> 为 true 时调用本函数。
template <typename T>
constexpr MPI_Datatype
mpi_datatype_of() noexcept
{
    if constexpr (std::is_same_v<T, char>)
        return MPI_CHAR;
    else if constexpr (std::is_same_v<T, signed char>)
        return MPI_SIGNED_CHAR;
    else if constexpr (std::is_same_v<T, unsigned char>)
        return MPI_UNSIGNED_CHAR;
    else if constexpr (std::is_same_v<T, short>)
        return MPI_SHORT;
    else if constexpr (std::is_same_v<T, unsigned short>)
        return MPI_UNSIGNED_SHORT;
    else if constexpr (std::is_same_v<T, int>)
        return MPI_INT;
    else if constexpr (std::is_same_v<T, unsigned int>)
        return MPI_UNSIGNED;
    else if constexpr (std::is_same_v<T, long>)
        return MPI_LONG;
    else if constexpr (std::is_same_v<T, unsigned long>)
        return MPI_UNSIGNED_LONG;
    else if constexpr (std::is_same_v<T, long long>)
        return MPI_LONG_LONG;
    else if constexpr (std::is_same_v<T, unsigned long long>)
        return MPI_UNSIGNED_LONG_LONG;
    else if constexpr (std::is_same_v<T, float>)
        return MPI_FLOAT;
    else if constexpr (std::is_same_v<T, double>)
        return MPI_DOUBLE;
    else if constexpr (std::is_same_v<T, long double>)
        return MPI_LONG_DOUBLE;
    else if constexpr (std::is_same_v<T, bool>)
        return MPI_C_BOOL;
    else if constexpr (std::is_same_v<T, std::byte>)
        return MPI_BYTE;
    else if constexpr (std::is_same_v<T, std::complex<float>>)
        return detail::ComplexDatatypeManager::instance().get_complex_float();
    else if constexpr (std::is_same_v<T, std::complex<double>>)
        return detail::ComplexDatatypeManager::instance().get_complex_double();
    else
        return MPI_DATATYPE_NULL;
}

// 在 MPI_Finalize 之前清理派生类型，避免在 finalize 之后释放导致未定义行为
inline void
cleanup_complex_types_before_finalize()
{
    detail::ComplexDatatypeManager::instance().cleanup_before_finalize();
}

} // namespace types
} // namespace compi
